from matplotlib import pyplot as plt
import cv2
import numpy as np
import tqdm
from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import pose as mp_pose
import poseembedding as pe                      # 姿态关键点编码模块
import poseclassifier as pc                     # 姿态分类器
import resultsmooth as rs                       # 分类结果平滑
import counter                                  # 动作计数器
import visualizer as vs                         # 可视化模块


def show_image(img, figsize=(10, 10)):
    """显示输出的 PIL 图像."""
    plt.figure(figsize=figsize)
    plt.imshow(img)
    plt.show()


def video_process(video_path, flag):
    # 指定视频路径和输出名称
    # video_path = 'pushup-sample.mp4'
    # class_name 需要与你的训练样本的两个动作状态图像文件夹的名字中的一个（或者是与 fitness_poses_csvs_out 中的一个 csv 文件的名字）保持一致，
    # 它后面将用于分类时的索引。具体是哪个动作文件夹的名字取决于你的运动是什么，例如：如果是深蹲，明显比较重要的判断计数动作是蹲下去；
    # 如果是俯卧撑，则判断计数的动作是向上拉到最高点的那个动作；如果是俯卧撑，则判断计数的动作是最低点的那个动作
    if flag == 1:
        class_name = 'push_down'
        out_video_path = 'pushup-sample-out.mp4'
    elif flag == 2:
        class_name = 'squat_down'
        out_video_path = 'squat-sample-out.mp4'


    # 打开视频
    video_cap = cv2.VideoCapture(video_path)

    # 获取一些视频参数以生成带分类的输出视频。
    video_n_frames = video_cap.get(cv2.CAP_PROP_FRAME_COUNT)
    video_fps = video_cap.get(cv2.CAP_PROP_FPS)
    video_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    video_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # 初始化跟踪器、分类器和计数器。在每个视频之前都要这样做，因为它们都有状态。

    # 包含姿势类 CSV 文件的文件夹。这应该与构建分类器时输出 CSV 的文件夹相同。
    pose_samples_folder = 'fitness_poses_csvs_out'

    # 初始化跟踪器。
    pose_tracker = mp_pose.Pose()

    # 初始化嵌入器。
    pose_embedder = pe.FullBodyPoseEmbedder()

    # 初始化分类器。
    # 确保你使用的参数与引导时相同。
    pose_classifier = pc.PoseClassifier(
        pose_samples_folder=pose_samples_folder,
        pose_embedder=pose_embedder,
        top_n_by_max_distance=30,
        top_n_by_mean_distance=10)

    # # 取消注释以验证分类器使用的目标姿势并找到异常值。
    # outliers = pose_classifier.find_pose_sample_outliers()
    # print('Number of pose sample outliers (consider removing them): ', len(outliers))

    # 初始化 EMA 平滑。
    pose_classification_filter = rs.EMADictSmoothing(
        window_size=10,
        alpha=0.2)

    # 初始化计数器。
    repetition_counter = counter.RepetitionCounter(
        class_name=class_name,
        enter_threshold=6,
        exit_threshold=4)

    # 初始化渲染器。
    pose_classification_visualizer = vs.PoseClassificationVisualizer(
        class_name=class_name,
        plot_x_max=video_n_frames,
        # 如果与 `top_n_by_mean_distance` 相同，图形看起来更好。
        plot_y_max=10)

    # 在视频上运行分类。

    # 打开输出视频。
    out_video = cv2.VideoWriter(out_video_path, cv2.VideoWriter_fourcc(*'mp4v'), video_fps, (video_width, video_height))

    frame_idx = 0
    output_frame = None
    with tqdm.tqdm(total=video_n_frames, position=0, leave=True) as pbar:
        while True:
            # 获取视频的下一帧。
            success, input_frame = video_cap.read()
            if not success:
                break

            # 运行姿势跟踪器。
            input_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2RGB)
            result = pose_tracker.process(image=input_frame)
            pose_landmarks = result.pose_landmarks

            # 绘制姿势预测。
            output_frame = input_frame.copy()
            if pose_landmarks is not None:
                mp_drawing.draw_landmarks(
                    image=output_frame,
                    landmark_list=pose_landmarks,
                    connections=mp_pose.POSE_CONNECTIONS)

            if pose_landmarks is not None:
                # 获取关键点。
                frame_height, frame_width = output_frame.shape[0], output_frame.shape[1]
                pose_landmarks = np.array([[lmk.x * frame_width, lmk.y * frame_height, lmk.z * frame_width]
                                           for lmk in pose_landmarks.landmark], dtype=np.float32)
                assert pose_landmarks.shape == (33, 3), 'Unexpected landmarks shape: {}'.format(pose_landmarks.shape)

                # 在当前帧上对姿势进行分类。
                pose_classification = pose_classifier(pose_landmarks)

                # 使用 EMA 进行平滑分类。
                pose_classification_filtered = pose_classification_filter(pose_classification)

                # 计数重复次数。
                repetitions_count = repetition_counter(pose_classification_filtered)
            else:
                # 没有姿势 => 当前帧没有分类。
                pose_classification = None

                # 仍然向过滤器中添加空分类以保持未来帧的正确平滑。
                pose_classification_filtered = pose_classification_filter(dict())
                pose_classification_filtered = None

                # 假设人 '冻结' 时不更新计数器。只需取最新的重复次数。
                repetitions_count = repetition_counter.n_repeats

            # 绘制分类图和重复计数器。
            output_frame = pose_classification_visualizer(
                frame=output_frame,
                pose_classification=pose_classification,
                pose_classification_filtered=pose_classification_filtered,
                repetitions_count=repetitions_count)

            # 保存输出帧。
            out_video.write(cv2.cvtColor(np.array(output_frame), cv2.COLOR_RGB2BGR))

            frame_idx += 1
            pbar.update()

    # 关闭输出视频。
    out_video.release()

    # 释放 MediaPipe 资源。
    pose_tracker.close()

    # 显示视频的最后一帧，并将其保存为图像文件
    if output_frame is not None:
        # 显示图像
        show_image(output_frame)
        # 将最后一帧保存为图像文件
        cv2.imwrite('static/last_frame.jpg', cv2.cvtColor(np.array(output_frame), cv2.COLOR_RGB2BGR))

